"""
向量数据库服务层

Author: wenbin
Create: 2024-09-12
"""
import os

import weaviate
from injector import inject
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStoreRetriever
from langchain_openai import OpenAIEmbeddings
from langchain_weaviate import WeaviateVectorStore


@inject
class VectorDatabaseService:
    """向量数据库服务"""

    def __init__(self):
        self.client = weaviate.connect_to_local(host=os.getenv("WEAVIATE_HOST"), port=int(os.getenv("WEAVIATE_PORT")))
        self.vector_store = WeaviateVectorStore(
            client=self.client,
            index_name="KnowledgeLibrary",
            text_key="text",
            embedding=OpenAIEmbeddings(model="text-embedding-3-small")
        )

    def get_retriever(self) -> VectorStoreRetriever:
        """获取检索器"""
        return self.vector_store.as_retriever(search_type="mmr")

    @staticmethod
    def format_documents(documents: list[Document]) -> str:
        """将文档列表格式化为字符串表达"""
        return "\n".join([document.page_content for document in documents])
